# 本程序 计划弄清dolfinx程序变量的值的操作
# 关注重点是 刚度矩阵和载荷向量
# Demo
# u(x, y) = sin(x)cos(y) f(x, y) = 2sin(x)cos(y)
# -Laplace(u) = f in \Omega
# u = g on \partial\Omega
# \Omega = [0, 5]\times[0, 1] 长方形 
import importlib.util

if importlib.util.find_spec("petsc4py") is not None:
    import dolfinx

    if not dolfinx.has_petsc:
        print("This demo requires DOLFINx to be compiled with PETSc enabled.")
        exit(0)
    from petsc4py.PETSc import ScalarType  # type: ignore
else:
    print("This demo requires petsc4py.")
    exit(0)

from mpi4py import MPI

import numpy as np
from petsc4py import PETSc
import ufl
from dolfinx import fem, io, mesh, plot
from dolfinx.fem.petsc import LinearProblem, assemble_matrix, assemble_vector, create_matrix
from ufl import ds, dx, grad, inner
from dolfinx.io import XDMFFile


# +
msh = mesh.create_rectangle(
    comm=MPI.COMM_WORLD,
    points=((0.0, 0.0), (5.0, 1.0)),
    n=(40, 8),
    cell_type=mesh.CellType.triangle,
)
V = fem.functionspace(msh, ("Lagrange", 1))


facets = mesh.locate_entities_boundary(
    msh,
    dim=(msh.topology.dim - 1),
    marker=lambda x: np.isclose(x[0], 0.0) | np.isclose(
        x[0], 5.0) | np.isclose(x[1], 0.0) | np.isclose(x[1], 1.0)
)
u = ufl.TrialFunction(V)
v = ufl.TestFunction(V)
x = ufl.SpatialCoordinate(msh)


def u_true(x):
    return np.sin(x[0]) * np.cos(x[1])


f = 2 * ufl.sin(x[0]) * ufl.cos(x[1])
a = inner(grad(u), grad(v)) * dx  # <class 'ufl.form.Form'>
a_form = fem.form(a)  # < class 'dolfinx.fem.forms.Form' >
L = inner(f, v) * dx
L_form = fem.form(L)

u_err = dolfinx.fem.Function(V)
u_exact = dolfinx.fem.Function(V)
u_exact.interpolate(u_true)

dofs = fem.locate_dofs_topological(V=V, entity_dim=1, entities=facets)
bc = fem.dirichletbc(u_exact, dofs=dofs)



problem = LinearProblem(a, L, bcs=[bc], petsc_options={
                        "ksp_type": "preonly", "pc_type": "lu"})
uh = problem.solve()
uh.name = "uh"  # 设置名字 画图更好区分

with XDMFFile(MPI.COMM_WORLD, f"demo-result/poisson.xdmf", "w") as ufile_xdmf:
    uh.x.scatter_forward()  # uh是数值结果 type= dolfinx.fem.Function
    ufile_xdmf.write_mesh(msh)
    ufile_xdmf.write_function(uh)
# 生成.h5文件和.xdmf文件
